package org.ctstudio

import org.apache.commons.math3.exception.InsufficientDataException
import org.apache.commons.math3.ml.clustering.Cluster
import org.apache.commons.math3.ml.clustering.Clusterable
import org.apache.commons.math3.ml.clustering.evaluation.ClusterEvaluator
import org.apache.commons.math3.util.MathArrays
import org.ctstudio.extensions.div
import org.ctstudio.extensions.plusAssign

/**
 * Compute the Calinski and Harabasz score.
 * It is also known as the Variance Ratio Criterion.
 * @author chentao106
 * @see <https://www.tandfonline.com/doi/abs/10.1080/03610927408827101>
 * java版正在合入apache commons math4的提案中。
 */
class CalinskiHarabasz<T : Clusterable> : ClusterEvaluator<T>(null) {
    /**
     * {@inheritDoc}
     */
    override fun score(clusters: List<Cluster<T>>): Double {
        val dimension = dimensionOfClusters(clusters)
        val centroid = meanOfClusters(clusters, dimension)
        var intraDistanceProduct = 0.0
        var extraDistanceProduct = 0.0
        for (cluster in clusters) {
            // Calculate the center of the cluster.
            val clusterCentroid = mean(cluster.points, dimension)
            for (p in cluster.points) {
                // Increase the intra distance sum
                intraDistanceProduct += covariance(clusterCentroid, p.point)
            }
            // Increase the extra distance sum
            extraDistanceProduct += cluster.points.size * covariance(centroid, clusterCentroid)
        }
        val pointCount = countAllPoints(clusters)
        val clusterCount = clusters.size
        // Return the ratio of the intraDistranceProduct to extraDistanceProduct
        return if (intraDistanceProduct == 0.0) 1.0 else extraDistanceProduct * (pointCount - clusterCount) /
                (intraDistanceProduct * (clusterCount - 1))
    }

    /**
     * Returns whether the first evaluation score is considered to be better
     * than the second one by this evaluator.
     *
     *
     * larger score is better.
     *
     * @param score1 the first score
     * @param score2 the second score
     * @return `true` if the first score is considered to be better, `false` otherwise
     */
    override fun isBetterScore(score1: Double, score2: Double): Boolean {
        return score1 > score2
    }

    /**
     * Calculate covariance of two double array.
     * <pre>
     * covariance = \sum\limits_{i=1}^n sigma_(p1\[i]-p2\[i])^2)
    </pre> *
     *
     * @param p1 Double array
     * @param p2 Double array
     * @return covariance of two double array
     */
    private fun covariance(p1: DoubleArray, p2: DoubleArray): Double {
        MathArrays.checkEqualLength(p1, p2)
        var sum = 0.0
        for (i in p1.indices) {
            val dp = p1[i] - p2[i]
            sum += dp * dp
        }
        return sum
    }

    /**
     * Calculate the mean of all the points.
     *
     * @param points    A collection of points
     * @param dimension The dimension of each point
     * @return The mean value.
     */
    private fun mean(points: Collection<T>, dimension: Int): DoubleArray {
        var centroid = DoubleArray(dimension)
        points.forEach { centroid.plusAssign(it.point) }
        centroid /= points.size.toDouble()
        return centroid
    }

    /**
     * Calculate the mean of all the points in the clusters.
     *
     * @param clusters  A collection of clusters.
     * @param dimension The dimension of each point.
     * @return The mean value.
     */
    private fun meanOfClusters(
        clusters: Collection<Cluster<T>?>,
        dimension: Int
    ): DoubleArray {
        var centroid = DoubleArray(dimension)
        var allPointsCount = 0
        for (cluster in clusters) {
            cluster!!.points.forEach {
                centroid.plusAssign(it.point)
                allPointsCount++
            }
        }
        centroid /= allPointsCount.toDouble()
        return centroid
    }

    /**
     * Count all the points in collection of cluster.
     *
     * @param clusters collection of cluster
     * @return points count
     */
    private fun countAllPoints(clusters: Collection<Cluster<T>>): Int {
        return clusters.sumBy { it.points.size }
    }

    /**
     * Detect the dimension of points in the clusters
     *
     * @param clusters collection of cluster
     * @return The dimension of the first point in clusters
     */
    private fun dimensionOfClusters(clusters: Collection<Cluster<T>>): Int {
        // Iteration and find out the first point.
        return clusters.firstOrNull()?.points?.firstOrNull()?.point?.size ?: throw InsufficientDataException()
    }
}